
import time
from src.face_align import norm_crop 
from src.utils import drawImg,is_point_near_segment
from src.global_dict import gd
import ast
import datetime
import cv2
from pathlib import Path

import heapq

def is_target_out_of_frame(bbox, frame_shape, margin=100):

    x1, y1, x2, y2 = map(int, bbox)
    # print(x1, y1, x2, y2)
    height, width = frame_shape[0], frame_shape[1]

  
    if (
        x1 < margin or            
        y1 < margin or            
        x2 > width - margin or    
        y2 > height - margin      
    ):
        return True
    return False
def update_top_faces_heap(result_details, track_id, new_face_conf, new_face_bgr ,max_faces=5):

    faces_heap = result_details[track_id]['face_bgr']

    if len(faces_heap) < max_faces:
   
        heapq.heappush(faces_heap, (new_face_conf, new_face_bgr))
    else:
       
        if new_face_conf > faces_heap[0][0]:
            heapq.heapreplace(faces_heap, (new_face_conf, new_face_bgr))

    result_details[track_id]['face_bgr'] = faces_heap

def ensure_path_exists(path_str):
    path = Path(path_str)
    path.mkdir(parents=True, exist_ok=True)
    return path

def post_process(img_bgr, task_id, args, logger, yolo_face, yolo11n_pose, rule_info_dict,frame_count,result_details,q_endt,track_id_list,end_track_id_list):

    try:
        end_line = rule_info_dict.get('end_line', [])
        end_line_list = ast.literal_eval(end_line)
    except (ValueError, SyntaxError) as e:
        logger.error(f"Error parsing end_line: {e}")
        return

    name_list = []

    # t1 = time.time()
    results = yolo11n_pose.track(img_bgr,imgsz=640, tracker="bytetrack.yaml", persist=True, device=0, iou=0.5, conf=0.25, classes=[0],verbose=False) 
    # results = yolo11n_pose(img_bgr,imgsz=640, device=0, iou=0.7, conf=0.25, verbose=False)  # predict on an image
    # t2 = time.time()
    # annotated_frame = results[0].plot()
    # cv2.line(annotated_frame, end_line_list[0], end_line_list[1], (255, 0, 0), 2)
    # cv2.imwrite('img/' + str(frame_count) + '.jpg', annotated_frame)
    # print(t2-t1)
    # print(frame_count)
    
    for result in results[0]:
        if result.boxes is None or result.boxes.id is None:
            continue
        xyxy = result.boxes.xyxy.cpu().numpy().tolist()[0]
        track_id = result.boxes.id.int().cpu().tolist()[0]
        # print(track_id)
        x1, y1, x2, y2 = map(int, xyxy)

        if abs(int(end_line_list[0][1]) - y2) <= 1000:
            try:
                img_bgr_body = img_bgr[y1:y2, x1:x2].copy()
                bboxes_face, kpss = yolo_face.detect(img_bgr_body, args.face_detec_thres2)
                if len(bboxes_face) == 0:
                    continue
                for bbox_face, kps in zip(bboxes_face, kpss):
                    face_conf = float(bbox_face[4])
                    face_bgr = norm_crop(img_bgr_body, kps)
                    update_top_faces_heap(result_details, track_id, face_conf, face_bgr)

                if abs(int(end_line_list[0][1]) - y2) <= 50 and track_id not in end_track_id_list:
                    formatted_end_t = datetime.datetime.fromtimestamp(time.time()).strftime('%Y_%m_%d_%H_%M_%S_%f')
                    result_details[track_id]['endt'] = formatted_end_t
                    cv2.line(img_bgr, end_line_list[0], end_line_list[1], (255, 0, 0), 2)
                    body_path = ensure_path_exists(f"zhuang_xian/{task_id}/")
                    cv2.imwrite(f"{body_path}/{track_id}_{formatted_end_t}.jpg", img_bgr)
                    end_track_id_list.append(track_id)

                person_leave = is_target_out_of_frame(xyxy, img_bgr.shape)
                if person_leave and track_id not in track_id_list and result_details[track_id]['face_bgr']:
                    q_endt.put({track_id: result_details[track_id]})
                    track_id_list.append(track_id)
                    if len(track_id_list) >= 10000:
                        track_id_list.pop(0)

            except Exception as e:
                logger.error(f"Error processing track_id {track_id}: {e}")

    t6 = time.time()
    # cost = t6 - gd.lastt_dict[task_id]
    # gd.lastt_dict[task_id] = t6
    # print("cost: ", cost)







        
    


